import os.path
import pathlib
from typing import Dict
from pprint import pprint
import sqlite3

from flask import Flask, url_for, redirect, request, send_from_directory, g

from pluginUtils.utils.loader import get_comprehensioner, get_reader
from pluginUtils.basic import BasicReaderPlugin, BasicComprehensionPlugin
from utils import load_config

from reader_api_file_blueprint import api_blueprint as file_api_blueprint
from reader_api_plugin_blueprint import api_blueprint as plugin_api_blueprint
from reader_api_db_blueprint import api_blueprint as db_api_blueprint, get_db, database_name

app = Flask(__name__)

ALLOWED_EXTENSIONS = list()
readertype2pluginobject: Dict[str, BasicReaderPlugin] = dict()
comprehensionername2pluginobject: Dict[str, BasicComprehensionPlugin] = dict()
last_reader_plugin = 'None'  # 用于实现模型内存管理
last_comprehension_plugin = 'None'


@app.route('/')
def hello_world():
    return redirect(url_for('static', filename='index.html'))


@app.route('/favicon.ico')
def favicon():
    return send_from_directory(os.path.join(app.root_path, 'static'),
                               'favicon.ico', mimetype='image/vnd.microsoft.icon')


@app.route("/test/run")
def run():
    code = request.args['code']
    value = exec(code)
    print(value)

    return ""


@app.teardown_appcontext
def close_connection(exception):
    db = getattr(g, '_database', None)
    if db is not None:
        db.close()


@app.route("/api/query", methods=['POST'])
def query():
    global last_reader_plugin, last_comprehension_plugin
    cur = get_db().cursor()

    contents = request.json
    bookname = contents['bookname']
    query_ = contents['query']
    comprehensionername = contents['plugincomprehensionername']
    explained = contents["explained"]
    num_retrieval = contents['num_retrieval']
    params = contents['params']
    params = {param['name']: param['default'] for param in params}

    assert comprehensionername in comprehensionername2pluginobject, "未知的阅读理解模型"
    suffix = bookname[bookname.rfind("."):]
    if suffix != last_reader_plugin:
        if last_reader_plugin != 'None':
            reader.hibernate()
        if reader.mainreader != suffix:
            reader.setmainreader(suffix)
        print("wake reader plugin of", suffix)
        reader.wake()
        print("-" * 60)
        last_reader_plugin = suffix
    if reader.get_readfile_name() != bookname or reader.blocked_in_read:
        if os.path.isfile(os.path.join("bookshelf", bookname)):
            print("reading", bookname)
            bookpath = os.path.join("bookshelf", bookname)
            reader.readbook(bookpath)
        else:
            return {
                "answer": "无效数据源"
            }
    destence, context = reader.search(query_, k=num_retrieval)
    destence, context = destence[0], context[0]
    if comprehensionername != last_comprehension_plugin:
        if last_comprehension_plugin != "None":
            print(last_comprehension_plugin, "hibernating")
            comprehensionername2pluginobject[last_comprehension_plugin].hibernate()
        print("wake comprehensioner plugin of", comprehensionername)
        comprehensionername2pluginobject[comprehensionername].wake()
        last_comprehension_plugin = comprehensionername
        print("-" * 60)
    answer = comprehensionername2pluginobject[comprehensionername].inference(context, query_, explained, **params)
    if explained:
        answer, explaination = answer
        res = {
            "context": context,
            "destence": destence.tolist(),
            "items": explaination.items
        }
    else:
        res = dict()
    res["answer"] = answer
    cur.execute("INSERT INTO usage VALUES (?, ?, ?)", (query_, answer, "".join(context)))
    get_db().commit()
    pprint(res)
    return res


if __name__ == '__main__':
    config = load_config()

    if not pathlib.Path(database_name).is_file():
        con_userdata = sqlite3.connect(database_name)
        cur_userdata = con_userdata.cursor()
        cur_userdata.execute("CREATE TABLE usage (question, answer, context)")
        cur_userdata.execute("CREATE TABLE revise (question, answer, context, revise)")
        cur_userdata.close()
        con_userdata.close()

    comprehensionername2pluginobject = get_comprehensioner(config)
    reader, readertype2pluginobject = get_reader(config)
    print("support types", reader.get_types())
    app.config['UPLOAD_FOLDER'] = 'bookshelf'
    app.config['reader'] = reader
    app.config['ALLOWED_EXTENSIONS'] = reader.get_types()
    app.config['config'] = config
    app.config['comprehensionername2pluginobject'] = comprehensionername2pluginobject
    app.config['readertype2pluginobject'] = readertype2pluginobject
    app.register_blueprint(file_api_blueprint, url_prefix="/api")
    app.register_blueprint(plugin_api_blueprint, url_prefix="/api")
    app.register_blueprint(db_api_blueprint, url_prefix="/api")
    app.run(host="0.0.0.0")
